from calculate import load_and_process_data  # 导入数据处理函数
from model import RadarCNN  # 导入模型
import torch
import torch.nn as nn
import torch.optim as optim
from torch.utils.data import DataLoader, TensorDataset

# 数据加载和处理
def prepare_dataloader(features, labels, batch_size=64):
    features = torch.tensor(features, dtype=torch.float32)
    labels = torch.tensor(labels, dtype=torch.float32)  # 回归任务使用float类型
    dataset = TensorDataset(features, labels)
    dataloader = DataLoader(dataset, batch_size=batch_size, shuffle=True)
    return dataloader

# 训练函数
# def train_model(model, dataloader, criterion, optimizer, num_epochs=10):
#     for epoch in range(num_epochs):
#         model.train()
#         running_loss = 0.0
#         for inputs, targets in dataloader:
#             optimizer.zero_grad()
#             outputs = model(inputs).squeeze()  # 输出维度调整为 (batch_size,)
#             loss = criterion(outputs, targets)
#             loss.backward()
#             optimizer.step()
#             running_loss += loss.item()
#         print(f"Epoch {epoch + 1}/{num_epochs}, Loss: {running_loss / len(dataloader)}")

def train_model(model, dataloader, criterion, optimizer, num_epochs=60):
    """
    训练模型并打印每个 epoch 的损失、相关系数和均方根误差。

    参数:
    - model: 要训练的模型
    - dataloader: 数据加载器
    - criterion: 损失函数
    - optimizer: 优化器
    - num_epochs: 训练轮数
    """
    best_loss = float('inf')
    for epoch in range(num_epochs):
        model.train()
        running_loss = 0.0
        all_targets = []
        all_predictions = []

        for inputs, targets in dataloader:
            optimizer.zero_grad()
            outputs = model(inputs).squeeze()  # 输出维度调整为 (batch_size,)
            loss = criterion(outputs, targets)
            loss.backward()
            optimizer.step()

            # 累计损失
            running_loss += loss.item()

            # 保存所有的预测值和真实值用于后续计算
            all_targets.append(targets.detach())
            all_predictions.append(outputs.detach())

        # 将所有批次的目标值和预测值拼接成一个张量
        all_targets = torch.cat(all_targets)
        all_predictions = torch.cat(all_predictions)

        # 计算相关系数
        target_mean = all_targets.mean()
        pred_mean = all_predictions.mean()
        covariance = ((all_targets - target_mean) * (all_predictions - pred_mean)).mean()
        target_std = all_targets.std()
        pred_std = all_predictions.std()
        correlation = covariance / (target_std * pred_std)

        # 计算均方根误差（RMSE）
        rmse = torch.sqrt(((all_targets - all_predictions) ** 2).mean())

        # 打印结果
        avg_loss = running_loss / len(dataloader)
        print(f"Epoch {epoch + 1}/{num_epochs}, Loss: {avg_loss:.4f}, RMSE: {rmse:.4f}, Correlation: {correlation:.4f}")

        if loss.item() < best_loss:
            best_loss = loss.item()
            torch.save(model.state_dict(), "best_model9.0.pth")

if __name__ == "__main__":
    from calculate import load_and_process_data  # 导入数据处理函数

    # 定义数据路径和文件名
    root_path = 'E:/福建数据/spar/04/test1/'
    filename_ZH = 'Z_RADR_I_Z9600_20240430024720_O_DOR_PAR-SD_CAP_025.bin.gz_ZH.csv'
    filename_CC = 'Z_RADR_I_Z9600_20240430024720_O_DOR_PAR-SD_CAP_025.bin.gz_CC.csv'

    # 加载和处理数据
    features, labels = load_and_process_data(root_path, filename_ZH, filename_CC)

    print(features.shape,labels.shape)

    # 调整形状为 (batch_size, channels, height, width)
    features = features.reshape(-1, 3, 3, 3)

    # 准备数据
    dataloader = prepare_dataloader(features, labels)

    # 初始化模型、损失函数和优化器
    model = RadarCNN()
    criterion = nn.MSELoss()  # 均方误差损失
    optimizer = torch.optim.Adam(model.parameters(), lr=0.001)

    # 训练模型
    train_model(model, dataloader, criterion, optimizer)


